import os
import shutil
from modelscope import snapshot_download
from tqdm import tqdm

class ModelDownloader:
    def __init__(self, cache_dir):
        self.cache_dir = cache_dir

    def check_disk_space(self, required_space=10*1024*1024*1024):  # 默认需要10GB空间
        """检查磁盘空间是否足够"""
        disk_usage = shutil.disk_usage(self.cache_dir)
        return disk_usage.free > required_space

    def download_model(self, model_id, output_dir=None):
        """
        下载指定ID的模型
        Args:
            model_id: ModelScope模型ID
            output_dir: 可选的输出目录
        Returns:
            str: 模型下载路径
        """
        if not self.check_disk_space():
            raise Exception("磁盘空间不足，请确保至少有10GB可用空间")

        try:
            model_dir = snapshot_download(
                model_id,
                cache_dir=output_dir or self.cache_dir,
                revision='master',
                ignore_file_pattern=['.git', '.gitattributes', '.gitignore']
            )
            return model_dir
        except Exception as e:
            raise Exception(f"模型下载失败: {str(e)}") 